package weiyao;

import weiyao.GUI.VulnerabilityRecord;

import java.sql.*;
import java.util.List;
import java.util.Map;

public class DatabaseUtil {
    // 数据库配置
    private static final String DB_URL = "jdbc:mysql://localhost:3306/member";
    private static final String USER = "root";
    private static final String PASS = "123456";

    public static void saveVulnerabilities(List<VulnerabilityRecord> records) throws SQLException {
        if (records == null || records.isEmpty()) return;

        try (Connection conn = DriverManager.getConnection(DB_URL, USER, PASS)) {
            conn.setAutoCommit(false); // 使用事务

            for (VulnerabilityRecord record : records) {
                switch (record.getType()) {
                    case "SQL注入":
                        saveSQLInjection(conn, record);
                        break;
                    case "XSS":
                        saveXSS(conn, record);
                        break;
                    case "CSRF":
                        saveCSRF(conn, record);
                        break;
                    case "弱口令":
                        saveWeakPassword(conn, record);
                        break;
                }
            }

            conn.commit();
        } catch (SQLException e) {
            throw new SQLException("数据库操作失败: " + e.getMessage(), e);
        }
    }

    private static void saveSQLInjection(Connection conn, VulnerabilityRecord record) throws SQLException {
        // 移除了 form_details 字段
        String sql = "INSERT INTO sql_injections (url, payloadandshow) VALUES (?, ?)";
        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            Map<String, String> details = record.getDetails();
            String payload = getValueOrNA(details, "载荷");
            pstmt.setString(1, record.getUrl());
            pstmt.setString(2, payload);
            pstmt.executeUpdate();
        }
    }

    private static void saveXSS(Connection conn, VulnerabilityRecord record) throws SQLException {
        // 移除了 context 字段
        String sql = "INSERT INTO xss_vulnerabilities (url, payload) VALUES (?, ?)";
        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            Map<String, String> details = record.getDetails();
            String payload = getValueOrNA(details, "载荷");
            pstmt.setString(1, record.getUrl());
            pstmt.setString(2, payload);
            pstmt.executeUpdate();
        }
    }

    private static void saveCSRF(Connection conn, VulnerabilityRecord record) throws SQLException {
        // 将 payload 改回 token_status
        String sql = "INSERT INTO csrf_vulnerabilities (url, form_action, token_status) VALUES (?, ?, ?)";
        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            Map<String, String> details = record.getDetails();
            String formAction = getValueOrNA(details, "URL");
            String tokenStatus = getValueOrNA(details, "载荷"); // 载荷表示令牌状态
            pstmt.setString(1, record.getUrl());
            pstmt.setString(2, formAction);
            pstmt.setString(3, tokenStatus);
            pstmt.executeUpdate();
        }
    }

    private static void saveWeakPassword(Connection conn, VulnerabilityRecord record) throws SQLException {
        String sql = "INSERT INTO weak_passwords (url, username, password) VALUES (?, ?, ?)";
        try (PreparedStatement pstmt = conn.prepareStatement(sql)) {
            Map<String, String> details = record.getDetails();
            String username = getValueOrNA(details, "用户名");
            String password = getValueOrNA(details, "密码");
            pstmt.setString(1, record.getUrl());
            pstmt.setString(2, username);
            pstmt.setString(3, password);
            pstmt.executeUpdate();
        }
    }

    // 辅助方法：获取值或返回"N/A"
    private static String getValueOrNA(Map<String, String> details, String key) {
        String value = details.get(key);
        return (value == null || value.isEmpty()) ? "N/A" : value;
    }
}